# 责任链模式 - 数据处理模块基类
import asyncio
import copy
import re
from typing import List, Dict

from backend.app.common.log import log
from backend.app.schemas.device import DeviceInfoSchema, TenderReq, TenderVectorSchema
from backend.app.schemas.device_v2 import DeviceSchemaV2, TenderDeviceReqV2
from backend.app.services.chain_service import summarize_text, get_vector_client
from backend.app.services.tender.search import TenderDeviceMatchScoreChain


async def tender_search(req: TenderReq):
    """
    标书匹配
    :param req: TenderReq
    :return:
    """
    # 格式化请求参数
    step_one = TenderReqFormatStep()
    question_list = await step_one.process(req) or []
    # 批量搜索问题答案
    pipline = AsyncDataPipeline()
    for question in question_list:
        pipline.add_module(TenderMatchStep, question, platform=req.platform, country_sort=req.country)
        # pipline.add_module(TenderSearchDeviceStep, question, shop_id=shop_id, country_sort=req.country)

    answer_list = await pipline.process_data()
    # 将所有答案匹配标书
    plan_step_2 = TenderBrandStep()
    plan_step_1 = TenderCatenaStep(next_module=plan_step_2)
    device_map = await plan_step_1.process(answer_list)

    # 填充请求数据
    output_step = TenderOutputStep()
    output = await output_step.process(device_map, req=req)
    return output


# 异步版本的数据处理模块基类
class AsyncDataProcessingModule:
    def __init__(self, next_module=None, is_async=False):
        self.next_module = next_module
        self.is_async = is_async

    async def process(self, data, **kwargs):
        processed_data = await self._process(data, **kwargs)
        if self.next_module:
            return await self.next_module.process(processed_data, **kwargs)
        return processed_data

    async def _process(self, data, **kwargs):
        raise NotImplementedError("Subclasses must implement _process method.")


class TenderSearchDeviceStep(AsyncDataProcessingModule):
    """
    单个问题查找匹配
    RAG强化搜索
    """

    @staticmethod
    def format_sources(source_documents: List, page_content: TenderVectorSchema):
        """
        将向量中的ID跟向量库搜索结果比对，并格式化
        :param source_documents:
        :param page_content: TenderVectorSchema
        :return: DeviceInfoSchema
        """
        child_list = []
        device = None
        choice_g_id = int(page_content.g_id) if page_content.g_id and page_content and page_content.g_id.strip().strip(
            ".").isdigit() else 0
        for s in source_documents:
            obj = DeviceInfoSchema(**s.metadata)
            is_same_gid = obj.gid == choice_g_id
            try:
                is_same_name = re.search(page_content.device_name, obj.device_name)
            except Exception as e:
                log.warning(f"[format_sources] re.search err: {e}")
                is_same_name = False

            if is_same_gid or (choice_g_id < 1 and is_same_name):
                device = obj
            else:
                child_list.append(obj)
        if not device:
            return
        device.child_list = child_list
        return device

    @staticmethod
    def get_country_desc(country_sort: int) -> str:
        """映射设备描述信息"""
        # 国别排序
        country_sort_desc = {1: "中国", 2: "中国以外的其他国家"}
        country = country_sort_desc.get(country_sort, "") or ""
        return country

    async def _process(self, question, **kwargs):
        """
        :param question: TenderDeviceReq
        :return: Dict[str: DeviceInfoSchema], {question_hash_id: DeviceInfoSchema}
        """
        shop_id = kwargs.get("shop_id")
        country_sort = kwargs.get("country_sort")
        result = dict(hash_id=question.hash_id, device=question)
        meter = question.tender_meter
        meter = await summarize_text(meter)
        country = self.get_country_desc(country_sort)
        client = get_vector_client()

        try:
            out_put = await client.execute(shop_id=shop_id, tender_name=question.device_name, country=country,
                                           attribute=meter)
            # log.info(f"[TenderSearchDeviceStep] out_put query: {out_put['query']}")
            log.info(f"[TenderSearchDeviceStep] out_put result: {out_put['result']}")
        except asyncio.TimeoutError as e:
            log.warning(f"[TenderSearchDeviceStep][_process] time out err: {e}")
            return result
        except Exception as e:
            log.warning(f"[TenderSearchDeviceStep][_process] err: {e}")
            return result

        # 向量数据库未查到结果   未匹配到如何处理
        if not (out_put['result'] and len(out_put["source_documents"])):
            log.info(
                f"[TenderSearchDeviceStep][_process] 没有查到相应的结果, source_documents: {len(out_put['source_documents'] or [])} \n")
            return result
        page_content = client.parser(out_put["result"])

        # page_content = self.match_result(out_put["result"])
        # 未检测匹配结果时将以下两种情况加入匹配列表（根据设备名称匹配(未匹配到)、默认获取第一个元数据）
        match_fields = [page_content]
        is_match = page_content.g_id != "" if page_content else False
        device_name = page_content.device_name if page_content else ""
        source_documents = out_put["source_documents"]
        if not is_match:
            log.info(
                f"[TenderSearchDeviceStep][_process] 未匹配到结果, source_documents: {['{}, {}'.format(obj.metadata['g_id'], obj.metadata['device_name']) for obj in out_put['source_documents']]} \n")

        if not is_match and device_name:
            match_fields.append(TenderVectorSchema(device_name=device_name, g_id=""))

        if not is_match and not source_documents:
            gid = int(source_documents[0].metadata["g_id"])
            match_fields.append(TenderVectorSchema(g_id=str(gid)))

        device_info = None
        for match_vector in match_fields:
            device_info = self.format_sources(out_put["source_documents"], match_vector)
            if device_info and device_info.gid > 0:
                break

        if not (device_info and device_info.gid):
            log.info(f"[TenderSearchDeviceStep][_process] format_sources 未匹配到结果, query: {out_put['query']} \n")
            log.info(f"[TenderSearchDeviceStep][_process] format_sources 未匹配到结果, result: {out_put['result']} \n")
            log.info(
                f"[TenderSearchDeviceStep][_process] format_sources 未匹配到结果, source_documents: {['{}, {}'.format(obj.metadata['g_id'], obj.metadata['device_name']) for obj in out_put['source_documents']]} \n")
        return dict(hash_id=question.hash_id(), device=device_info)


class TenderMatchStep(AsyncDataProcessingModule):
    """标书匹配"""

    async def _process(self, question, **kwargs):
        """
        :param question: TenderDeviceReq
        :return: Dict[str: DeviceSchemaV2], {question_hash_id: DeviceSchemaV2}
        """

        query = "{}; {}".format(question.device_name, question.tender_meter)
        platform = kwargs.get("platform")

        result = dict(hash_id=question.hash_id(), device=question)
        try:
            query = await summarize_text(query)
            # 按照相关性排序
            # chain = TenderDeviceMatchChain(shop_id=shop_id)
            # 按照向量距离评分排序
            chain = TenderDeviceMatchScoreChain(platform)
            device_info = await chain.search(query)
            log.info(f"[TenderMatchStep][_process] query: {query}")
        except asyncio.TimeoutError as e:
            log.warning(f"[TenderMatchStep][_process] time out err: {e}")
            return result
        result["device"] = device_info
        return result


class TenderBrandStep(AsyncDataProcessingModule):
    """标书方案之品牌筛选"""

    @staticmethod
    def sync_device_brand(class_device: List[Dict]):
        """
        同步每个二级分类的品牌
        :param class_device: [{"hash_id": xxx, "device": DeviceSchemaV2}]
        :return: [{"hash_id": xxx, "device": DeviceSchemaV2}]
        """
        # 找到所有商品中共有的品牌 === (统计商品下所有品牌出现的次数，选择最大的一个)
        device_key = "device"
        if not class_device or len(class_device) <= 1:
            return class_device
        # 默认二级分类下第一个设备的品牌为选中品牌
        choice_brand = class_device[0][device_key].brand_id
        if not choice_brand:
            #  如果是空则跳过筛选
            return class_device

        for index, goods in enumerate(class_device):
            choice_item = goods[device_key]

            # 原来的品牌就符合
            if choice_item.brand_id == choice_brand:
                continue
            choice_child_list = choice_item.child_list
            for k, child in enumerate(choice_child_list):
                if child.brand_id == choice_brand:
                    # 替换来选中的元素
                    child_list = copy.deepcopy(choice_item.child_list)
                    choice_item.child_list = []
                    del child_list[k]
                    child_list.append(choice_item)
                    child.child_list = child_list
                    class_device[index][device_key] = child
                    log.info(
                        f"[TenderBrandStep][sync_device_brand] 品牌一致切换： choice_item: {choice_item.gid}，更换的设备ID：{child.gid}")
                    break
        return class_device

    async def _process(self, class_group: Dict, **kwargs) -> Dict:
        """
        :param class_group: {class_id: [{device=DeviceSchemaV2, hash_id=xxx}]}
        :return: {hash_id:DeviceSchemaV2}
        """
        device_map = dict()
        for class_id, class_device in class_group.items():
            device_list = class_device or []
            if class_id and len(class_device) > 1:
                device_list = self.sync_device_brand(class_device or [])

            for device in device_list:
                hash_id = device["hash_id"]
                device_map[hash_id] = device["device"]
        return device_map


class TenderCatenaStep(AsyncDataProcessingModule):
    """标书方案之系列筛选"""

    @staticmethod
    def get_class_info(data) -> Dict:
        """
        根据二级分类将设备分组
        :param data: [{device=DeviceInfoSchema, hash_id=xxx}]
        :return: {class_id: [{device=DeviceInfoSchema, hash_id=xxx}]}
        """
        class_group = dict()
        for device_info in data:
            device = device_info["device"]
            if not (device and device.gid):
                log.warning(f"[TenderCatenaStep][get_class_info] 未找到设备： device : {device}")
                continue
            class_id = device.class_id or 0
            if not class_group.get(class_id):
                class_group[class_id] = []
            class_group[class_id].append(device_info)
        return class_group

    @staticmethod
    def sync_device_catena(class_device: List[Dict]):
        """
        匹配同系列设备(去除非同系列的选项)
        :param class_device: [{device=DeviceSchemaV2, hash_id=xxx}]
        :return: [{"hash_id": xxx, "device": DeviceSchemaV2}]
        """
        device_key = "device"

        if not class_device or len(class_device) <= 1:
            return class_device
        # 默认第一个设备的系列为最终系列
        choice_catena = class_device[0][device_key].catena
        if not choice_catena:
            return class_device

        for index, goods in enumerate(class_device):
            choice_item = goods[device_key]
            # 原来的品牌就符合
            if choice_item.catena == choice_catena:
                continue

            for k, child in enumerate(choice_item.child_list):
                if child.catena == choice_catena:
                    # 替换来选中的元素
                    child_list = choice_item.child_list
                    choice_item.child_list = []
                    del child_list[k]
                    child_list.append(choice_item)
                    # 只留下符合系列的设备
                    child.child_list = [x for x in child_list if x.catena == choice_catena]
                    class_device[index][device_key] = child
                    log.info(
                        f"[TenderCatenaStep][sync_device_catena] 系列一致切换： choice_item: {choice_item.gid}，更换的设备ID：{child.gid}")
                    break
        return class_device

    async def _process(self, data: List, **kwargs):
        """
        :param data: [{device=DeviceInfoSchema, hash_id=xxx}]
        :return: {class_id: [{device=DeviceInfoSchema, hash_id=xxx}]}
        """
        class_group = self.get_class_info(data)
        for class_id, class_device in class_group.items():
            if not class_id:
                continue
            class_group[class_id] = self.sync_device_catena(class_device or [])
        return class_group


class TenderReqFormatStep(AsyncDataProcessingModule):
    """标书请求参数优化"""

    @staticmethod
    def get_device_list(req: TenderReq) -> List:
        """
        获取标书要求列表
        :param req: TenderReq
        :return: List[TenderDeviceReq]
        """
        device_list = []
        data = req.device_list
        for system in data:
            if not system.goods:
                continue
            for module in system.goods:
                if not module.device_list:
                    continue
                for device in module.device_list:
                    if not device.tender_meter:
                        log.info(f"[TenderReqFormat][get_device_map] 忽略参数：device: {device}")
                        continue
                    device.tender_name = device.device_name
                    device_list.append(device)
        return device_list

    async def _process(self, data: TenderReq, **kwargs):
        """
        格式化输入参数
        :param data: List[TenderSystemReq]
        :return: List[TenderDeviceReq]
        """
        device_list = self.get_device_list(data)
        return device_list


class TenderOutputStep(AsyncDataProcessingModule):
    """
    标书 格式化输出数据
    """

    @staticmethod
    def update_device(device: TenderDeviceReqV2, match_device: DeviceSchemaV2 | None = None, price_field: str = None):
        """
        更新设备字段
        :param price_field: 价格字段
        :param device: 请求参数  TenderDeviceReqV2
        :param match_device: 匹配到的设备： DeviceSchemaV2
        :return:
        """
        device.device_name = device.tender_name
        number = device.number
        if not (number and str(number).strip() != ""):
            device.number = 1
        if not match_device:
            return device

        device_info = device.model_dump(by_alias=True)
        match_device.device_name = device.tender_name
        # 更新相关字段
        child_list = match_device.child_list or []
        child_info_list = []
        for index, child in enumerate(child_list):
            device_item = copy.deepcopy(device_info)
            if price_field:
                device_item["priceList"][price_field] = child.price
            child.device_name = device_item["tender_name"]
            child_info = child.model_dump(by_alias=True)
            device_item.update(child_info)
            child_info_list.append(device_item)

        if price_field:
            device_info["priceList"][price_field] = match_device.price
        match_dict = match_device.model_dump(by_alias=True)
        match_dict["child_list"] = child_info_list

        device_info.update(match_dict)
        number = device_info.get("number")
        if not (number and str(number).strip() != ""):
            device_info["number"] = 1
        return TenderDeviceReqV2(**device_info)

    async def _process(self, data, **kwargs):
        """
        格式化输出数据
        :param data: {hash_id: DeviceInfoSchema}
        :param kwargs: {req: TenderReq}
        :return: TenderReq
        """
        req = kwargs["req"]
        price_field = req.price_name or "price"
        system_device = req.device_list
        for system_index, system in enumerate(system_device):
            if not system.goods:
                continue
            module_list = system.goods
            for module_index, module in enumerate(module_list):
                if not module.device_list:
                    continue
                device_list = module.device_list
                for device_index, device in enumerate(device_list):
                    hash_id = device.hash_id()
                    match_device = data.get(hash_id, None)
                    result = self.update_device(device, match_device=match_device, price_field=price_field)
                    price = float(result.price) if result.price else 0
                    req.device_list[system_index].price += price
                    req.device_list[system_index].goods[module_index].device_list[device_index] = result
        return req


# 异步版本的数据处理流水线
class AsyncDataPipeline:
    def __init__(self):
        self.modules = []

    def add_module(self, module, data, **kwargs):
        self.modules.append([module, data, kwargs])

    async def process_data(self):
        tasks = [module().process(value, **kwargs) for module, value, kwargs in self.modules]
        return await asyncio.gather(*tasks)
